# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/10 19:47 
@Author  : ZhangShenao 
@File    : main.py 
@Desc    : 主函数
"""
import logging

import dotenv
from flask import Flask, request, render_template

from embedding import VectorStore
from loading import DocumentLoader
from retrieval import QARetriever
from splitting import DocumentSplitter

# 加载环境变量
dotenv.load_dotenv()

# 加载文档列表
loader = DocumentLoader('doc')
docs = loader.load_docs()

# 将文档分割成分块
splitter = DocumentSplitter(chunk_size=200, chunk_overlap=10)
chunked_docs = splitter.split_documents(docs)

# 将分块文档存储到向量数据库中
vector_store = VectorStore.new_store_by_documents(chunked_docs)

# 创建Retriever检索器
retriever_chain = QARetriever.new_retriever_chain(vector_store)

# 设置Logging
logging.basicConfig()
logging.getLogger('langchain.retrievers.multi_query').setLevel(logging.INFO)

# 创建Flask Web服务
app = Flask(__name__)


# 绑定路由
@app.route('/', methods=['GET', 'POST'])
def home():
    if request.method == 'POST':
        # 接收用户输入作为问题
        question = request.form.get('question')

        # RetrievalQA链 - 读入问题，生成答案
        result = retriever_chain({"query": question})

        # 把大模型的回答结果返回网页进行渲染
        return render_template('index.html', result=result)

    return render_template('index.html')


if __name__ == '__main__':
    # 启动Flask服务
    app.run(host='0.0.0.0', debug=True, port=5000)
